use pdata::internal::data::protogen::collector::trace::v1::ExportTraceServiceRequest;
use pdata::{
    internal::data::protogen::collector::metrics::v1::ExportMetricsServiceRequest,
    protogen::collector::logs::v1::ExportLogsServiceRequest,
};
use serde::{Deserialize, Serialize};
use std::fmt::Display;
use std::path::{Path, PathBuf};
use tokio::fs::{self, OpenOptions};
use tokio::io::{AsyncReadExt, AsyncSeekExt, AsyncWriteExt};

/// Extension used for WAL files
const WAL_EXT: &str = ".wal";

/// Extension used for state files
const STATE_EXT: &str = ".state";

/// Temporary extension used for atomic writes/renames
const TEMP_EXT: &str = ".temp";

pub const TRACES_FILE: &str = "traces";
pub const METRICS_FILE: &str = "metrics";
pub const LOGS_FILE: &str = "logs";

#[derive(Clone)]
pub struct WalWriter {
    root_dir: PathBuf,
}

impl WalWriter {
    /// Create a new WAL writer
    pub fn new(root_dir: PathBuf) -> Self {
        Self { root_dir }
    }

    /// Append protobuf message to appropriate WAL file with automatic type detection
    pub async fn append<T>(&self, data: &T)
    where
        T: prost::Message + WalDataTypeDetector,
    {
        let (file_path, type_name) = data.wal_info(&self.root_dir);

        let result = async {
            let mut payload = Vec::new();
            prost::Message::encode(data, &mut payload)
                .map_err(|e| std::io::Error::new(std::io::ErrorKind::InvalidData, e))?;

            self.ensure_parent_dir(&file_path).await?;

            let mut f = OpenOptions::new()
                .create(true)
                .append(true)
                .open(&file_path)
                .await?;
            let len = payload.len() as u32;
            // Write header (length) and payload in a single write to avoid partial frames being observed by readers
            let mut frame = Vec::with_capacity(4 + payload.len());
            frame.extend_from_slice(&len.to_le_bytes());
            frame.extend_from_slice(&payload);
            f.write_all(&frame).await?;
            f.sync_data().await
        }
        .await;

        if let Err(e) = result {
            tracing::warn!(target: "wal", "failed to append {} to WAL: {:?}", type_name, e);
        }
    }

    /// Update sent state for a specific data type
    pub async fn update_sent_state(&self, data_type: WalDataType, sent_index: u64) -> bool {
        let state_path = data_type.state_path(&self.root_dir);
        let type_name = match data_type {
            WalDataType::Traces => "traces",
            WalDataType::Metrics => "metrics",
            WalDataType::Logs => "logs",
        };

        let temp_path =
            state_path.with_extension(format!("{}{}", STATE_EXT.trim_start_matches('.'), TEMP_EXT));
        let state = StateFile { sent_index };

        let result = async {
            self.ensure_parent_dir(&state_path).await?;
            let mut f = OpenOptions::new()
                .create(true)
                .write(true)
                .truncate(true)
                .open(&temp_path)
                .await?;
            let json = serde_json::to_vec(&state).unwrap_or_else(|_| b"{}".to_vec());
            f.write_all(&json).await?;
            f.flush().await?;
            f.sync_data().await.ok();
            fs::rename(&temp_path, &state_path).await
        }
        .await;

        if let Err(e) = result {
            tracing::warn!(target: "wal", "failed to update {} sent state: {:?}", type_name, e);
            false
        } else {
            true
        }
    }

    /// Read state from specific data type file
    pub async fn read_state(&self, data_type: WalDataType) -> Result<StateFile, WalError> {
        let state_path = data_type.state_path(&self.root_dir);
        match fs::read(&state_path).await {
            Ok(bytes) => match serde_json::from_slice::<StateFile>(&bytes) {
                Ok(state) => Ok(state),
                Err(e) => {
                    tracing::error!(target: "wal", "failed to parse state file as JSON: path={} error={}", state_path.display(), e);
                    Err(WalError::Json(e))
                }
            },
            Err(e) if e.kind() == std::io::ErrorKind::NotFound => Ok(StateFile::default()),
            Err(e) => {
                tracing::error!(target: "wal", "failed to read state file: path={} error={}", state_path.display(), e);
                Err(WalError::Io(e))
            }
        }
    }

    /// Compact the WAL file by removing frames up to the current sent index.
    ///
    /// Behavior:
    /// - Reads `sent_index` from the state file for `data_type`.
    /// - If `sent_index` >= current WAL size: truncate WAL to empty and reset state to 0.
    /// - Else: copy bytes from `sent_index..end` into a temporary WAL, fsync, then atomically rename temp over original.
    /// - Finally, reset state `sent_index` to 0 (data moved to start).
    pub async fn compact(&self, data_type: WalDataType) -> Result<(), WalError> {
        let wal = data_type.wal_path(&self.root_dir);
        if !wal.exists() {
            return Ok(());
        }

        let sent = self
            .read_state(data_type)
            .await
            .unwrap_or_default()
            .sent_index as usize;

        let bytes = fs::read(&wal).await.map_err(WalError::Io)?;
        if sent >= bytes.len() {
            let _ = OpenOptions::new()
                .write(true)
                .truncate(true)
                .open(&wal)
                .await
                .map_err(WalError::Io)?;
            let _ = self.update_sent_state(data_type, 0).await;
            return Ok(());
        }

        let tmp = wal.with_extension(format!(
            "{}{}",
            WAL_EXT.trim_start_matches('.'),
            TEMP_EXT
        ));
        self.ensure_parent_dir(&tmp).await.map_err(WalError::Io)?;

        let mut f = OpenOptions::new()
            .create(true)
            .write(true)
            .truncate(true)
            .open(&tmp)
            .await
            .map_err(WalError::Io)?;
        f.write_all(&bytes[sent..]).await.map_err(WalError::Io)?;
        f.sync_data().await.ok();
        fs::rename(&tmp, &wal).await.map_err(WalError::Io)?;
        let _ = self.update_sent_state(data_type, 0).await;
        Ok(())
    }

    /// Read frames and invoke an async callback per frame without using a channel
    pub async fn read_frames_streaming<F, Fut>(
        &self,
        data_type: WalDataType,
        start_index: u64,
        mut on_frame: F,
    ) -> std::io::Result<()>
    where
        F: FnMut(Vec<u8>) -> Fut,
        Fut: std::future::Future<Output = ()>,
    {
        let wal_path = data_type.wal_path(&self.root_dir);
        if !wal_path.exists() {
            return Ok(());
        }

        // Determine file size boundary at the start to avoid reading into the data being written
        let file_metadata = fs::metadata(&wal_path).await?;
        let read_end_boundary = file_metadata.len();

        if start_index >= read_end_boundary {
            return Ok(()); // No new data to read
        }

        let mut f = OpenOptions::new().read(true).open(&wal_path).await?;
        f.seek(std::io::SeekFrom::Start(start_index)).await?;
        let mut current_pos = start_index;
        let mut buf4 = [0u8; 4];

        loop {
            // Check if there is enough space to read the frame header (4 bytes)
            if current_pos + 4 > read_end_boundary {
                break; // Avoid reading beyond the boundary
            }

            match f.read_exact(&mut buf4).await {
                Ok(_) => {
                    let len = u32::from_le_bytes(buf4) as usize;
                    current_pos += 4;

                    // Check if there is enough space to read the complete payload
                    if current_pos + len as u64 > read_end_boundary {
                        // This frame crosses the boundary, possibly incomplete, stop reading
                        tracing::warn!(target: "wal", "Frame at position {} may be incomplete, stopping read at boundary {}", current_pos - 4, read_end_boundary);
                        break;
                    }

                    let mut payload = vec![0u8; len];
                    f.read_exact(&mut payload).await?;
                    current_pos += len as u64;

                    on_frame(payload).await;
                }
                Err(e) => {
                    if e.kind() == std::io::ErrorKind::UnexpectedEof {
                        break;
                    } else {
                        return Err(e);
                    }
                }
            }
        }
        Ok(())
    }

    async fn ensure_parent_dir(&self, file_path: &Path) -> std::io::Result<()> {
        if let Some(parent) = file_path.parent() {
            if !parent.exists() {
                fs::create_dir_all(parent).await?;
            }
        }
        Ok(())
    }
}

/// Trait for automatic WAL data type detection
pub trait WalDataTypeDetector {
    fn wal_info(&self, root_dir: &Path) -> (PathBuf, &'static str);
}

#[derive(Debug, Clone, Copy)]
pub enum WalDataType {
    Traces,
    Metrics,
    Logs,
}

impl Display for WalDataType {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        let name = match self {
            WalDataType::Traces => "traces",
            WalDataType::Metrics => "metrics",
            WalDataType::Logs => "logs",
        };
        write!(f, "{}", name)
    }
}

impl WalDataType {
    /// Get the full WAL file path for this data type
    pub fn wal_path(self, root: &Path) -> PathBuf {
        let file_name = match self {
            WalDataType::Traces => TRACES_FILE,
            WalDataType::Metrics => METRICS_FILE,
            WalDataType::Logs => LOGS_FILE,
        };
        root.join(format!("{}{}", file_name, WAL_EXT))
    }

    /// Get the full state file path for this data type
    pub fn state_path(self, root: &Path) -> PathBuf {
        let file_name = match self {
            WalDataType::Traces => TRACES_FILE,
            WalDataType::Metrics => METRICS_FILE,
            WalDataType::Logs => LOGS_FILE,
        };
        root.join(format!("{}{}", file_name, STATE_EXT))
    }
}

#[derive(Debug)]
pub enum WalError {
    Io(std::io::Error),
    Json(serde_json::Error),
}

impl std::fmt::Display for WalError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            WalError::Io(e) => write!(f, "io error: {}", e),
            WalError::Json(e) => write!(f, "json error: {}", e),
        }
    }
}

impl std::error::Error for WalError {}

#[derive(Debug, Serialize, Deserialize, Clone, Default)]
pub struct StateFile {
    #[serde(default)]
    pub sent_index: u64,
}

impl WalDataTypeDetector for ExportTraceServiceRequest {
    fn wal_info(&self, root_dir: &Path) -> (PathBuf, &'static str) {
        (
            root_dir.join(format!("{}{}", TRACES_FILE, WAL_EXT)),
            "traces",
        )
    }
}

impl WalDataTypeDetector for ExportMetricsServiceRequest {
    fn wal_info(&self, root_dir: &Path) -> (PathBuf, &'static str) {
        (
            root_dir.join(format!("{}{}", METRICS_FILE, WAL_EXT)),
            "metrics",
        )
    }
}

impl WalDataTypeDetector for ExportLogsServiceRequest {
    fn wal_info(&self, root_dir: &Path) -> (PathBuf, &'static str) {
        (root_dir.join(format!("{}{}", LOGS_FILE, WAL_EXT)), "logs")
    }
}
